
#include "register/op_def_registry.h"
#include "tiling_common.h"
#include "triu_tiling.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext *context) {

  TriuTilingData tiling;
  context->SetBlockDim(1);
  auto xshape = context->GetInputShape(0)->GetOriginShape();
  auto w = xshape.GetDim(xshape.GetDimNum() - 2);
  tiling.set_w(w);
  auto h = xshape.GetDim(xshape.GetDimNum() - 1);
  tiling.set_h(h);
  int64_t b = xshape.GetShapeSize() / w / h;
  tiling.set_b(b);

  auto diagonal = context->GetAttrs()->GetInt(0);
  tiling.set_diagonal(*diagonal);

  COMMON_TILING_GET_PLATFORM;
  COMMON_TILING_GET_UB_SIZE;

  auto dtype = context->GetInputDesc(0)->GetDataType();
  auto dtypeSize = GetSizeByDataType(dtype);

  // x * 2 + y * 2
  auto split_ub = int32_t(float(ub_size) / (4.0f)) / 256 * 256 / dtypeSize;
  tiling.set_tileLength(split_ub);
  tiling_zero_workspace(context);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  std::cout << "b = " << b << std::endl;
  std::cout << "w = " << tiling.get_w() << std::endl;
  std::cout << "h = " << tiling.get_h() << std::endl;
  std::cout << "diagonal = " << *diagonal << std::endl;
  std::cout << "tileLength = " << tiling.get_tileLength() << std::endl;

  return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext *context) {
  const gert::Shape *x1_shape = context->GetInputShape(0);
  gert::Shape *y_shape = context->GetOutputShape(0);
  *y_shape = *x1_shape;
  return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class Triu : public OpDef {
public:
  explicit Triu(const char *name) : OpDef(name) {
    this->Input("x")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Output("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Attr("diagonal").AttrType(OPTIONAL).Int(0);

    this->SetInferShape(ge::InferShape);

    this->AICore().SetTiling(optiling::TilingFunc);
    this->AICore().AddConfig("ascend310b");
  }
};

OP_ADD(Triu);
} // namespace ops
